/***************************************************************************
 *   Copyright (C) 2010 by programmeboy                                    *
 *   programmeboy@gmail.com                                                *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU General Public License as published by  *
 *   the Free Software Foundation; either version 2 of the License, or     *
 *   (at your option) any later version.                                   *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU General Public License     *
 *   along with this program; if not, write to the                         *
 *   Free Software Foundation, Inc.                                        *
 *   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             * 
 ***************************************************************************/

#include <ntifs.h>

#include "LDasm.h"
#include "ssdt.h"

extern PKSERVICE_TABLE_DESCRIPTOR KeServiceDescriptorTable;


NTSTATUS
QuerySSDTInfo(
    OUT PVOID Buffer,
    IN ULONG Length,
    OUT PULONG RetLength
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    if (!Buffer || Length < sizeof(KSERVICE_TABLE_DESCRIPTOR)){
        return STATUS_INVALID_PARAMETER;
    }
    
    //
    //  Just copy that
    //

    RtlCopyMemory(Buffer, KeServiceDescriptorTable, 
                  sizeof(KSERVICE_TABLE_DESCRIPTOR));
    
    if (RetLength){
        *RetLength = sizeof(KSERVICE_TABLE_DESCRIPTOR);
    }
    
    return STATUS_SUCCESS;
    
}

extern
NTKERNELAPI
BOOLEAN
KeAddSystemServiceTable(
    IN PULONG_PTR Base,
    IN PULONG Count OPTIONAL,
    IN ULONG Limit,
    IN PUCHAR Number,
    IN ULONG Index
);

PKSERVICE_TABLE_DESCRIPTOR 
QueryShadowSSDTAddress(
    VOID
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    /*ULONG Offset;
    ULONG Thread;
    NTSTATUS status;
    ULONG i, j;
    
    Offset = 0;
    for (i = 8; i < 0x8000; i += 4)
    {
        status = PsLookupThreadByThreadId((HANDLE)i, &(PETHREAD)Thread);
        if (NT_SUCCESS(status))
        {
            ObDereferenceObject((PETHREAD)Thread);
            if (!Offset)
            {
                
                //
                //  first find the offset of KTHREAD.ServiceTable
                //
                
                for (j = Thread; j < Thread + 0x1000; j += 4)
                {
                    
                    //
                    //  is not gui thread
                    //
                    
                    if (*(PULONG)j == (ULONG)&KeServiceDescriptorTable)
                    {
                        
                        //
                        //  get offset.
                        //
                        
                        Offset = j - Thread;
                        break;
                    }
                }
            }
            else
            {
                
                //
                //  yeah it's gui thread.
                //
                
                if (*(PULONG)(Thread + Offset) != (ULONG)&KeServiceDescriptorTable) 
                    return (PKSERVICE_TABLE_DESCRIPTOR)(*(PULONG)(Thread + Offset));
            }
        }
    }
    return 0;*/

    PUCHAR cPtr, pOpcode;
    ULONG Length;

    for (cPtr = (PUCHAR)KeAddSystemServiceTable;
         cPtr < (PUCHAR)KeAddSystemServiceTable + PAGE_SIZE;
         cPtr += Length){

        if (!MmIsAddressValid(cPtr)) break;

        Length = SizeOfCode(cPtr, &pOpcode);

        if (!Length || (Length == 1 && *pOpcode == 0xC3)) break;

        if (*(PUSHORT)pOpcode == 0x888D)
        {
            return  (PKSERVICE_TABLE_DESCRIPTOR)(*(PVOID *)(pOpcode + 2));
        }
    }
    return 0;
}

NTSTATUS
QueryShadowSSDTInfo(
    OUT PVOID Buffer,
    IN ULONG Length,
    OUT PULONG RetLength
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    PKSERVICE_TABLE_DESCRIPTOR KeShadowServiceDescriptorTable;

    if (!Buffer || Length < sizeof(KSERVICE_TABLE_DESCRIPTOR)){
        return STATUS_INVALID_PARAMETER;
    }
    
    KeShadowServiceDescriptorTable = QueryShadowSSDTAddress();
    if (!KeShadowServiceDescriptorTable){
        return STATUS_UNSUCCESSFUL;
    }
    

    //
    //  Just copy that
    //

    RtlCopyMemory(Buffer, KeShadowServiceDescriptorTable+1, 
        sizeof(KSERVICE_TABLE_DESCRIPTOR));

    if (RetLength){
        *RetLength = sizeof(KSERVICE_TABLE_DESCRIPTOR);
    }

    return STATUS_SUCCESS;
}

NTSTATUS
QueryShadowSSDTEntrys(
    OUT PVOID Buffer,
    IN ULONG Length,
    OUT PULONG RetLength
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    ULONG EntryLength;
    PKSERVICE_TABLE_DESCRIPTOR KeShadowServiceDescriptorTable;

    if (!Buffer){
        return STATUS_INVALID_PARAMETER;
    }
    
    KeShadowServiceDescriptorTable = QueryShadowSSDTAddress();
    if (!KeShadowServiceDescriptorTable){
        return STATUS_UNSUCCESSFUL;
    }
    
    KeShadowServiceDescriptorTable++;

    //
    //  calculate the size of ssdt entry
    //
    

    EntryLength = KeShadowServiceDescriptorTable->Limit * sizeof(ULONG);
    if (Length < EntryLength){
        return STATUS_BUFFER_TOO_SMALL;
    }
    
    RtlCopyMemory(Buffer, KeShadowServiceDescriptorTable->Base, EntryLength);
    
    if (RetLength){
        *RetLength = EntryLength;
    }
    return STATUS_SUCCESS;
}

NTSTATUS
QuerySSDTEntrys(
    OUT PVOID Buffer,
    IN ULONG Length,
    OUT PULONG RetLength
)
/*++

Routine Description:

    

Arguments:

	

Return Value:

--*/
{
    ULONG EntryLength;

    if (!Buffer){
        return STATUS_INVALID_PARAMETER;
    }
    
    //
    //  calculate the size of ssdt entry
    //
    
    EntryLength = KeServiceDescriptorTable->Limit * sizeof(ULONG);
    if (Length < EntryLength){
        return STATUS_BUFFER_TOO_SMALL;
    }
    
    RtlCopyMemory(Buffer, KeServiceDescriptorTable->Base, EntryLength);
    
    if (RetLength){
        *RetLength = EntryLength;
    }
    return STATUS_SUCCESS;
}
